package com.luckzj.sec.bridge;

import io.netty.channel.Channel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;

public class BridgeMapper {
    private static Logger logger = LoggerFactory.getLogger(BridgeMapper.class);

    private Map<Channel, Channel> s2d = new HashMap<>();
    private Map<Channel, Channel> d2s = new HashMap<>();

    public void addMap(Channel serverChannel, Channel clientChannel) {
        if (s2d.containsKey(serverChannel) || d2s.containsKey(clientChannel)) {
            if (logger.isWarnEnabled()) {
                logger.warn("duplicate src channel or dst channel detected. src port = {}, dst port = {}"
                        , serverChannel.localAddress(), clientChannel.remoteAddress());
            }
        }

        s2d.put(serverChannel, clientChannel);
        d2s.put(clientChannel, serverChannel);
    }

    public Channel getDst(Channel ch) {
        return s2d.get(ch);
    }

    public Channel getSrc(Channel ch) {
        return d2s.get(ch);
    }

    public void removeSrc(Channel ch) {
        Channel dst = s2d.get(ch);
        s2d.remove(ch);

        if (dst != null) {
            d2s.remove(dst);
        }
    }

    public void removeDst(Channel ch) {
        Channel dst = d2s.get(ch);
        d2s.remove(ch);

        if (dst != null) {
            s2d.remove(dst);
        }
    }
}
